{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b72c437f",
   "metadata": {
    "papermill": {
     "duration": 17.608072,
     "end_time": "2024-05-12T09:32:54.628939",
     "exception": false,
     "start_time": "2024-05-12T09:32:37.020867",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\ProgramData\\miniconda3\\envs\\study_dl\\Lib\\site-packages\\torchvision\\io\\image.py:13: UserWarning: Failed to load image Python extension: '[WinError 127] 找不到指定的程序。'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start on cpu device.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "import torchvision\n",
    "from PIL import Image\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchvision import models, transforms\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print(\"start on {} device.\".format(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd7b9b4b",
   "metadata": {
    "papermill": {
     "duration": 0.012772,
     "end_time": "2024-05-12T09:32:54.645141",
     "exception": false,
     "start_time": "2024-05-12T09:32:54.632369",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 保存训练模型名称\n",
    "save_model_name = \"vgg16_base_cifar10\"\n",
    "# 训练的轮数\n",
    "epoch = 50\n",
    "# 一批数据大小\n",
    "batch_size = 128\n",
    "# 丢弃率\n",
    "dropout = 0.5\n",
    "# 数据集路径\n",
    "# /kaggle/input/dogs-vs-cats\n",
    "dataset_path = \"./dataset\"\n",
    "# tensorboard路径\n",
    "writer = SummaryWriter(\"./logs\")\n",
    "# 学习率\n",
    "# learning_rate = 0.01\n",
    "# 1e-2=1 x (10)^(-2) = 1 /100 = 0.01\n",
    "learning_rate = 1e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "626c4c23",
   "metadata": {
    "papermill": {
     "duration": 7.188299,
     "end_time": "2024-05-12T09:33:01.836545",
     "exception": false,
     "start_time": "2024-05-12T09:32:54.648246",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "训练数据集长度: 50000\n",
      "测试数据集长度: 10000\n",
      "torch.Size([128, 3, 32, 32])\n",
      "torch.Size([128, 3, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.ToTensor()\n",
    "train_data = torchvision.datasets.CIFAR10(dataset_path, train=True, download=True,\n",
    "                                          transform=transform)\n",
    "train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=False)\n",
    "\n",
    "test_data = torchvision.datasets.CIFAR10(dataset_path, train=False, download=True,\n",
    "                                         transform=transform)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=0, drop_last=False)\n",
    "\n",
    "train_data_size = len(train_data)\n",
    "test_data_size = len(test_data)\n",
    "print(\"训练数据集长度: {}\".format(train_data_size))\n",
    "print(\"测试数据集长度: {}\".format(test_data_size))\n",
    "\n",
    "for data in train_dataloader:\n",
    "    imgs, targets = data\n",
    "    print(imgs.shape)\n",
    "    break\n",
    "\n",
    "for data in test_dataloader:\n",
    "    imgs, targets = data\n",
    "    print(imgs.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c9387983",
   "metadata": {
    "papermill": {
     "duration": 0.03027,
     "end_time": "2024-05-12T09:33:01.874180",
     "exception": false,
     "start_time": "2024-05-12T09:33:01.843910",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class VGG16_BASE(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VGG16_BASE, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(64, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(128, 128, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(256, 256, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(256, 256, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(256, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            # vgg16原本时 nn.AdaptiveAvgPool2d((7, 7))，由于输入图片大小导致\n",
    "            # nn.AvgPool2d(kernel_size=1, stride=1)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            # 原始模型vgg16输入image大小是224 x 224\n",
    "            # CIFAR10输入image大小是32 x 32\n",
    "            # 大小是小了 7 x 7倍\n",
    "            nn.Linear(512, 4096),\n",
    "            nn.ReLU(True),\n",
    "            nn.Dropout(p=dropout),\n",
    "            nn.Linear(4096, 4096),\n",
    "            nn.ReLU(True),\n",
    "            nn.Dropout(p=dropout),\n",
    "            nn.Linear(4096, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = self.features(x)\n",
    "        # output = output.view(output.size(0), -1)\n",
    "        output = torch.flatten(output, 1)\n",
    "        output = self.classifier(output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7b5cd18a",
   "metadata": {
    "papermill": {
     "duration": 0.560507,
     "end_time": "2024-05-12T09:33:02.441802",
     "exception": false,
     "start_time": "2024-05-12T09:33:01.881295",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VGG16_BASE(\n",
      "  (features): Sequential(\n",
      "    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU(inplace=True)\n",
      "    (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (5): ReLU(inplace=True)\n",
      "    (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (8): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (9): ReLU(inplace=True)\n",
      "    (10): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (11): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (12): ReLU(inplace=True)\n",
      "    (13): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (14): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (15): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (16): ReLU(inplace=True)\n",
      "    (17): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (18): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (19): ReLU(inplace=True)\n",
      "    (20): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (21): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (22): ReLU(inplace=True)\n",
      "    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (24): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (25): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (26): ReLU(inplace=True)\n",
      "    (27): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (28): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (29): ReLU(inplace=True)\n",
      "    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (31): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (32): ReLU(inplace=True)\n",
      "    (33): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (35): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (36): ReLU(inplace=True)\n",
      "    (37): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (38): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (39): ReLU(inplace=True)\n",
      "    (40): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (41): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (42): ReLU(inplace=True)\n",
      "    (43): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (classifier): Sequential(\n",
      "    (0): Linear(in_features=512, out_features=4096, bias=True)\n",
      "    (1): ReLU(inplace=True)\n",
      "    (2): Dropout(p=0.5, inplace=False)\n",
      "    (3): Linear(in_features=4096, out_features=4096, bias=True)\n",
      "    (4): ReLU(inplace=True)\n",
      "    (5): Dropout(p=0.5, inplace=False)\n",
      "    (6): Linear(in_features=4096, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# 网络模型相关初始化\n",
    "# 创建网络模型\n",
    "mymod = VGG16_BASE()\n",
    "print(mymod)\n",
    "\n",
    "mymod = mymod.to(device)\n",
    "# 损失函数\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "# loss_fn = loss_fn.to(device)\n",
    "\n",
    "# 优化器\n",
    "optimizer = torch.optim.SGD(mymod.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "47744705-96b5-4178-a2b6-3084aea6800e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 64, 32, 32]           1,792\n",
      "       BatchNorm2d-2           [-1, 64, 32, 32]             128\n",
      "              ReLU-3           [-1, 64, 32, 32]               0\n",
      "            Conv2d-4           [-1, 64, 32, 32]          36,928\n",
      "       BatchNorm2d-5           [-1, 64, 32, 32]             128\n",
      "              ReLU-6           [-1, 64, 32, 32]               0\n",
      "         MaxPool2d-7           [-1, 64, 16, 16]               0\n",
      "            Conv2d-8          [-1, 128, 16, 16]          73,856\n",
      "       BatchNorm2d-9          [-1, 128, 16, 16]             256\n",
      "             ReLU-10          [-1, 128, 16, 16]               0\n",
      "           Conv2d-11          [-1, 128, 16, 16]         147,584\n",
      "      BatchNorm2d-12          [-1, 128, 16, 16]             256\n",
      "             ReLU-13          [-1, 128, 16, 16]               0\n",
      "        MaxPool2d-14            [-1, 128, 8, 8]               0\n",
      "           Conv2d-15            [-1, 256, 8, 8]         295,168\n",
      "      BatchNorm2d-16            [-1, 256, 8, 8]             512\n",
      "             ReLU-17            [-1, 256, 8, 8]               0\n",
      "           Conv2d-18            [-1, 256, 8, 8]         590,080\n",
      "      BatchNorm2d-19            [-1, 256, 8, 8]             512\n",
      "             ReLU-20            [-1, 256, 8, 8]               0\n",
      "           Conv2d-21            [-1, 256, 8, 8]         590,080\n",
      "      BatchNorm2d-22            [-1, 256, 8, 8]             512\n",
      "             ReLU-23            [-1, 256, 8, 8]               0\n",
      "        MaxPool2d-24            [-1, 256, 4, 4]               0\n",
      "           Conv2d-25            [-1, 512, 4, 4]       1,180,160\n",
      "      BatchNorm2d-26            [-1, 512, 4, 4]           1,024\n",
      "             ReLU-27            [-1, 512, 4, 4]               0\n",
      "           Conv2d-28            [-1, 512, 4, 4]       2,359,808\n",
      "      BatchNorm2d-29            [-1, 512, 4, 4]           1,024\n",
      "             ReLU-30            [-1, 512, 4, 4]               0\n",
      "           Conv2d-31            [-1, 512, 4, 4]       2,359,808\n",
      "      BatchNorm2d-32            [-1, 512, 4, 4]           1,024\n",
      "             ReLU-33            [-1, 512, 4, 4]               0\n",
      "        MaxPool2d-34            [-1, 512, 2, 2]               0\n",
      "           Conv2d-35            [-1, 512, 2, 2]       2,359,808\n",
      "      BatchNorm2d-36            [-1, 512, 2, 2]           1,024\n",
      "             ReLU-37            [-1, 512, 2, 2]               0\n",
      "           Conv2d-38            [-1, 512, 2, 2]       2,359,808\n",
      "      BatchNorm2d-39            [-1, 512, 2, 2]           1,024\n",
      "             ReLU-40            [-1, 512, 2, 2]               0\n",
      "           Conv2d-41            [-1, 512, 2, 2]       2,359,808\n",
      "      BatchNorm2d-42            [-1, 512, 2, 2]           1,024\n",
      "             ReLU-43            [-1, 512, 2, 2]               0\n",
      "        MaxPool2d-44            [-1, 512, 1, 1]               0\n",
      "           Linear-45                 [-1, 4096]       2,101,248\n",
      "             ReLU-46                 [-1, 4096]               0\n",
      "          Dropout-47                 [-1, 4096]               0\n",
      "           Linear-48                 [-1, 4096]      16,781,312\n",
      "             ReLU-49                 [-1, 4096]               0\n",
      "          Dropout-50                 [-1, 4096]               0\n",
      "           Linear-51                   [-1, 10]          40,970\n",
      "================================================================\n",
      "Total params: 33,646,666\n",
      "Trainable params: 33,646,666\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.01\n",
      "Forward/backward pass size (MB): 6.75\n",
      "Params size (MB): 128.35\n",
      "Estimated Total Size (MB): 135.12\n",
      "----------------------------------------------------------------\n",
      "tensor([3, 5, 8, 5, 0, 8, 9, 7, 1, 1, 8, 5, 5, 2, 8, 9, 6, 0, 8, 1, 6, 0, 4, 5,\n",
      "        0, 8, 3, 8, 4, 8, 9, 6, 7, 9, 8, 5, 0, 3, 3, 8, 3, 9, 5, 5, 4, 8, 9, 2,\n",
      "        3, 7, 7, 4, 6, 6, 6, 7, 2, 6, 0, 8, 4, 4, 9, 8, 2, 9, 5, 2, 2, 0, 0, 8,\n",
      "        6, 0, 3, 0, 9, 1, 8, 0, 3, 5, 0, 5, 7, 5, 3, 9, 0, 4, 2, 2, 2, 8, 9, 4,\n",
      "        5, 9, 5, 7, 9, 6, 2, 7, 3, 0, 3, 6, 7, 2, 0, 1, 4, 4, 5, 3, 2, 0, 2, 2,\n",
      "        5, 3, 5, 1, 3, 1, 9, 8])\n",
      "torch.Size([128, 3, 32, 32])\n",
      "torch.Size([128, 3, 32, 32])\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'VGG16_BASE' object has no attribute 'model'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 12\u001b[0m\n\u001b[0;32m     10\u001b[0m x, targets \u001b[38;5;241m=\u001b[39m data\n\u001b[0;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(x\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m---> 12\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(mymod\u001b[38;5;241m.\u001b[39mmodel)):\n\u001b[0;32m     13\u001b[0m     x \u001b[38;5;241m=\u001b[39m mymod\u001b[38;5;241m.\u001b[39mmodel[i](x)\n\u001b[0;32m     14\u001b[0m     \u001b[38;5;28mprint\u001b[39m(i, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m layer shape：\u001b[39m\u001b[38;5;124m'\u001b[39m, x\u001b[38;5;241m.\u001b[39mshape)\n",
      "File \u001b[1;32mC:\\ProgramData\\miniconda3\\envs\\study_dl\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1688\u001b[0m, in \u001b[0;36mModule.__getattr__\u001b[1;34m(self, name)\u001b[0m\n\u001b[0;32m   1686\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01min\u001b[39;00m modules:\n\u001b[0;32m   1687\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m modules[name]\n\u001b[1;32m-> 1688\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m object has no attribute \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'VGG16_BASE' object has no attribute 'model'"
     ]
    }
   ],
   "source": [
    "input = torch.ones((64, 3, 32, 32))\n",
    "writer.add_graph(mymod, input)\n",
    "from torchsummary import summary\n",
    "summary(mymod, (3, 32, 32))\n",
    "for data in train_dataloader:\n",
    "    imgs, targets = data\n",
    "    print(targets)\n",
    "    print(imgs.shape)\n",
    "    # 打印每层的结果，方便对比哪层出现问题\n",
    "    x, targets = data\n",
    "    print(x.shape)\n",
    "    for i in range(len(mymod.model)):\n",
    "        x = mymod.model[i](x)\n",
    "        print(i, ' layer shape：', x.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e43e0708",
   "metadata": {
    "papermill": {
     "duration": 813.263212,
     "end_time": "2024-05-12T09:46:35.712576",
     "exception": false,
     "start_time": "2024-05-12T09:33:02.449364",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 保存训练数据绘图\n",
    "train_loss_list = []\n",
    "train_accuracy_list = []\n",
    "\n",
    "test_loss_list = []\n",
    "test_accuracy_list = []\n",
    "\n",
    "epoch_step = []\n",
    "print(\"start train.\")\n",
    "for i in range(epoch):\n",
    "    epoch_step.append(i)\n",
    "    \n",
    "    mymod.train()\n",
    "    start_time = time.time()\n",
    "    total_train_loss_pre_epoch = 0\n",
    "    total_train_accuracy_pre_epoch = 0\n",
    "    for data in train_dataloader:\n",
    "        imgs, targets = data\n",
    "        # print(imgs.shape)\n",
    "        imgs = imgs.to(device)\n",
    "        targets = targets.to(device)\n",
    "        outputs = mymod(imgs)\n",
    "        accuracy = (outputs.argmax(1) == targets).sum()\n",
    "        total_train_accuracy_pre_epoch += accuracy.item()\n",
    "\n",
    "        loss = loss_fn(outputs, targets)\n",
    "        total_train_loss_pre_epoch += loss.item() * targets.size(0)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    end_train_time = time.time()\n",
    "    print(\"第{}轮训练用时：{} s\".format(i, end_train_time - start_time))\n",
    "    print(\"第{}轮训练loss：{}\".format(i, total_train_loss_pre_epoch / train_data_size))\n",
    "    print(\"第{}轮训练精度：{}\".format(i, total_train_accuracy_pre_epoch / train_data_size))\n",
    "    print(\"------\")\n",
    "    train_loss_list.append(total_train_loss_pre_epoch / train_data_size)\n",
    "    train_accuracy_list.append(total_train_accuracy_pre_epoch / train_data_size)\n",
    "    \n",
    "    writer.add_scalar(\"train_loss\", total_train_loss_pre_epoch / train_data_size, i)\n",
    "    writer.add_scalar(\"train_acc\", total_train_accuracy_pre_epoch / train_data_size, i)\n",
    "    \n",
    "    \n",
    "    mymod.eval()\n",
    "    total_test_loss_pre_epoch = 0\n",
    "    total_test_accuracy_pre_epoch = 0\n",
    "    # 去掉梯度\n",
    "    with torch.no_grad():\n",
    "        for data in test_dataloader:\n",
    "            imgs, targets = data\n",
    "            imgs = imgs.to(device)\n",
    "            targets = targets.to(device)\n",
    "            outputs = mymod(imgs)\n",
    "\n",
    "            loss = loss_fn(outputs, targets)\n",
    "            total_test_loss_pre_epoch += loss.item() * targets.size(0)\n",
    "            # 统计在测试集上正确的个数，然后累加起来\n",
    "            accuracy = (outputs.argmax(1) == targets).sum()\n",
    "            total_test_accuracy_pre_epoch += accuracy.item()\n",
    "\n",
    "\n",
    "    end_test_time = time.time()\n",
    "    print(\"第{}轮测试用时：{} s\".format(i, end_test_time - end_train_time))\n",
    "    print(\"第{}轮测试loss：{}\".format(i, total_test_loss_pre_epoch / test_data_size))\n",
    "    print(\"第{}轮测试精度：{}\".format(i, total_test_accuracy_pre_epoch / test_data_size))\n",
    "    print(\"------\")\n",
    "    test_loss_list.append(total_test_loss_pre_epoch / test_data_size)\n",
    "    test_accuracy_list.append(total_test_accuracy_pre_epoch / test_data_size)\n",
    "    \n",
    "    writer.add_scalar(\"test_loss\", total_test_loss_pre_epoch / test_data_size, i)\n",
    "    writer.add_scalar(\"test_acc\", total_test_accuracy_pre_epoch / test_data_size, i)\n",
    "\n",
    "\n",
    "torch.save(mymod, \"./{}.pth\".format(save_model_name))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c181c554",
   "metadata": {
    "papermill": {
     "duration": 1.544573,
     "end_time": "2024-05-12T09:46:37.271348",
     "exception": false,
     "start_time": "2024-05-12T09:46:35.726775",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# plt.plot(x,y)\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, train_loss_list)\n",
    "plt.title(\"train_loss\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, test_loss_list)\n",
    "plt.title(\"test_loss\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, train_accuracy_list)\n",
    "plt.title(\"train_accuracy\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"acc\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, test_accuracy_list)\n",
    "plt.title(\"test_accuracy\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"acc\")\n",
    "plt.grid(True)\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, train_loss_list, label=\"train\")\n",
    "plt.plot(epoch_step, test_loss_list, label=\"test\")\n",
    "plt.legend()\n",
    "plt.title(\"loss\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(epoch_step, train_accuracy_list, label=\"train\")\n",
    "plt.plot(epoch_step, test_accuracy_list, label=\"test\")\n",
    "plt.legend()\n",
    "plt.title(\"accuracy\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"acc\")\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "databundleVersionId": 31148,
     "sourceId": 3362,
     "sourceType": "competition"
    }
   ],
   "dockerImageVersionId": 30699,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 845.542841,
   "end_time": "2024-05-12T09:46:39.813933",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-05-12T09:32:34.271092",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
